import operator
import os
import random
# 修复：添加 time 库用于短暂暂停，优化体验
import time
from typing import List, Dict, TypedDict, Annotated

from dotenv import load_dotenv
from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END

# 加载环境变量
load_dotenv()

# ---------------- 配置区域 ----------------
try:
    llm = ChatOpenAI(
        model=os.getenv("LONGMAO_MODEL"),
        base_url=os.getenv("LONGMAO_BASE_URL"),
        api_key=os.getenv("LONGMAO_API_KEY"),
        temperature=0.7
    )
except Exception as e:
    print(f"Error initializing LLM: {e}")
    exit(1)

# 真人玩家的名字
HUMAN_NAME = "我(真人)"

from sklearn.metrics.pairwise import cosine_similarity
from langchain_openai import OpenAIEmbeddings

# 初始化 Embedding 模型
try:
    embeddings_model = OpenAIEmbeddings(
        model=os.getenv("GUIJI_EMB_MODEL"),
        base_url=os.getenv("GUIJI_EMB_URL"),
        api_key=os.getenv("GUIJI_API_KEY")
    )
except Exception as e:
    print(f"Error initializing Embeddings: {e}")
    exit(1)


def calculate_similarity(text1, text2):
    """计算两个文本的余弦相似度"""
    if not text1 or not text2:
        return 0.0
    try:
        vec1 = embeddings_model.embed_query(text1)
        vec2 = embeddings_model.embed_query(text2)
        similarity = cosine_similarity([vec1], [vec2])[0][0]
        return similarity
    except Exception as e:
        print(f"Warning: Embedding calculation failed: {e}")
        return 0.0


# ---------------- 1. 定义游戏状态 ----------------
class GameState(TypedDict):
    players: List[str]
    player_words: Dict[str, str]
    spy_player: str
    history: Annotated[List[str], operator.add]
    current_turn_idx: int
    round_count: int
    votes: Dict[str, str]
    game_result: str
    cheater_detected: str
    # 新增：存储每一轮的相似度得分
    similarity_scores: Dict[str, float]


# ---------------- 2. 核心节点逻辑 ----------------

def init_game_node(state: GameState):
    """上帝节点：发牌"""
    print("\n" + "=" * 40)
    print("🎮 游戏开始！正在分配词汇...")
    word_pairs = [
        ("苹果", "梨"), ("微信", "QQ"), ("麦当劳", "肯德基"),
        ("成吉思汗", "努尔哈赤"), ("蜘蛛侠", "蝙蝠侠"), ("程序员", "黑客")
    ]
    common_word, spy_word = random.choice(word_pairs)
    players = [HUMAN_NAME, "AI_1号", "AI_2号"]
    random.shuffle(players)
    spy = random.choice(players)
    allocation = {}
    for p in players:
        allocation[p] = spy_word if p == spy else common_word

    human_word = allocation[HUMAN_NAME]
    identity = '卧底' if spy == HUMAN_NAME else '平民'
    print(f"🤫 嘘！你的身份是【{identity}】")
    print(f"🔑 你的词是：【{human_word}】 (切记：描述时不能直接包含这个词！)")
    print("-" * 40)
    print(f"👥 玩家发言顺序：{' -> '.join(players)}")
    print("=" * 40 + "\n")
    time.sleep(1)
    return {
        "players": players,
        "player_words": allocation,
        "spy_player": spy,
        "history": [],
        "current_turn_idx": 0,
        "round_count": 1,
        "votes": {},
        "cheater_detected": None,
        "similarity_scores": {}
    }


def speaker_node(state: GameState):
    """发言节点"""
    current_player = state["players"][state["current_turn_idx"]]
    my_word = state["player_words"][current_player]
    spy_player = state["spy_player"]
    round_num = state["round_count"]
    history_str = "\n".join(state['history']) if state['history'] else "（暂无，你是首个发言）"

    print(f"👉 轮到 [{current_player}] 发言...")

    if current_player == HUMAN_NAME:
        print(f"   (提示: 你的词是【{my_word}】)")
        while True:
            content = input(f"🎤 请输入你的描述 > ").strip()
            if not content: continue
            if my_word in content:
                print(f"🚫 犯规警告！不能直接包含关键词。请重试。")
                continue
            break
        msg = f"[{current_player}]: {content}"
    else:
        is_spy = (current_player == spy_player)
        if is_spy:
            prompt = f"""
                        你是{current_player}，身份【卧底】，词【{my_word}】。目标：隐藏自己。
                        第 {round_num} 轮。历史：{history_str}
                        思维链：
                        1. 风险评估：我是首发吗？前面人说了什么？
                        2. 伪装策略：描述【{my_word}】真实特征，但严禁使用其独有、精确的特征。使用宽泛、模糊或上位概念。
                        输出：仅一句话描述。
                        """
        else:
            prompt = f"""
            你是{current_player}，身份【平民】，词【{my_word}】。目标：找卧底。
            第 {round_num} 轮。历史：{history_str}
            思维链：
            1. 分析局势：谁的描述勉强、模糊或有冲突？
            2. 发言策略：描述一个相对核心特征向队友示好，但不要过于直白防止被猜到。
            输出：仅一句话描述。
            """
        try:
            # print(f"   (🤖 {current_player} 正在思考策略...)")
            response = llm.invoke([HumanMessage(content=prompt)])
            content = response.content.replace(my_word, "***").strip()
            if "最终输出" in content: content = content.split("最终输出")[-1].strip(": \n")
            msg = f"[{current_player}]: {content}"
            print(f"🤖 {msg}")
            time.sleep(random.uniform(0.5, 1.5))
        except Exception as e:
            msg = f"[{current_player}]: (大脑过载...跳过)"
            print(f"❌ AI生成失败: {e}")
    return {"history": [msg]}


def fact_check_node(state: GameState):
    """裁判节点"""
    current_player = state["players"][state["current_turn_idx"]]
    last_msg = state["history"][-1]
    real_word = state["player_words"][current_player]
    spoken_text = last_msg.split(":", 1)[1].strip() if ":" in last_msg else last_msg

    prompt = f"""
    裁判任务。词：【{real_word}】。描述："{spoken_text}"
    判断是否违规（FAIL）：1.撒谎/事实错误。2.直接包含词汇或过于直白导致游戏失去意义。
    模糊描述属于正常（PASS）。
    只回答：PASS 或 FAIL
    """
    try:
        check_res = llm.invoke([HumanMessage(content=prompt)]).content.strip().upper()
        if "FAIL" in check_res:
            print(f"\n🚨 裁判红牌！[{current_player}] 严重违规！直接淘汰！\n")
            return {"cheater_detected": current_player}
    except:
        pass
    return {"current_turn_idx": state["current_turn_idx"] + 1}


# ==================== 【重点】新增的语义分析节点 ====================
def semantic_analysis_node(state: GameState):
    """语义分析节点：计算每个人发言与平民词的相似度"""
    # 只有在所有人都说完话后才会进入这里
    print("\n🧠 === 正在进行语义向量分析(客观数据) ===")
    # 获取这一轮的发言（最后 N 条，N=玩家数）
    num_players = len(state["players"])
    current_turn_history = state["history"][-num_players:]

    # 获取平民词作为标尺
    spy_player = state["spy_player"]
    # 简单逻辑：如果真人不是卧底，真人的词就是平民词；否则 AI_1 的词是平民词
    common_word = state["player_words"][HUMAN_NAME] if spy_player != HUMAN_NAME else state["player_words"]["AI_1号"]
    print(f"   (经过判定，本局平民词标尺为：【{common_word}】)")

    scores = {}
    for msg in current_turn_history:
        try:
            player_name, content = msg.split(":", 1)
            player_name = player_name.strip("[]")
            content = content.strip()
            # 计算相似度
            score = calculate_similarity(common_word, content)
            scores[player_name] = round(score, 3)
            print(f"   📊 [{player_name}] 发言与【{common_word}】的语义相似度: {score:.3f}")
        except Exception as e:
            print(f"   ⚠️ 分析 [{msg}] 时出错: {e}")

    print("=================================\n")
    # 重要：将分析结果存入 state，传递给投票节点
    return {"similarity_scores": scores}


# ================================================================


def voting_node(state: GameState):
    """投票节点"""
    print(f"\n🗳️ === 第 {state['round_count']} 轮描述结束，开始投票 ===")
    print("-" * 30)
    players = state["players"]
    spy_player = state["spy_player"]
    history_str = "\n".join(state["history"])

    # 【重点】从 state 中取出刚才计算好的分数
    scores = state.get("similarity_scores", {})

    new_votes = {}
    for p in players:
        if p == HUMAN_NAME:
            # ... (真人投票逻辑不变) ...
            print("💀 请根据上面的历史发言，找出卧底。")
            print(f"   候选人: {players}")
            while True:
                vote_target = input(f"🗳️ 你投给谁? (输入名字的关键词) > ").strip()
                found = False
                for valid in players:
                    if vote_target and (valid in vote_target or vote_target in valid):
                        vote_target = valid
                        found = True
                        break
                if found: break
                print("❌ 找不到该玩家。")
            print(f"✅ 你投给了: {vote_target}")
        else:
            my_word = state["player_words"][p]
            is_spy = (p == spy_player)
            if is_spy:
                # ... (卧底投票逻辑不变) ...
                prompt = f"""
                 你是{p}，身份【卧底】，词【{my_word}】。目标：找替罪羊。
                 历史：{history_str}
                 思维链：1.确定对你威胁最大的人（描述最具体的平民）。2.定制谎言攻击他（攻击他太刻意、背书）。
                 输出：`被投人名字|简短理由`
                 """
            else:
                # ================= 🧑‍🤝‍🧑 平民投票策略 (集成语义数据) =================
                prompt = f"""
                你是 {p}，身份【平民】，词【{my_word}】。目标：找卧底。
                历史：{history_str}

                **【重要数据参考】本轮发言与平民词的语义相似度得分：**
                {scores}
                (分数越低，说明描述越模糊、信息量越少，嫌疑越大)

                思维链：
                1. **数据分析**：查看得分。谁的分数显著低于其他人？低分者通常是在进行极度模糊的“划水”，嫌疑极大。
                2. **综合判断**：优先怀疑得分最低的“划水怪”。如果没有明显低分项，再结合语感判断。

                输出要求：只输出 `被投人名字|简短理由(可引用分数)`
                """

            # AI 执行推理和投票
            try:
                # print(f"   (🤖 {p} 正在结合历史和数据进行推理...)")
                res = llm.invoke([HumanMessage(content=prompt)])
                content = res.content.strip()
                reason = "无理由"
                vote_target = content
                if "|" in content:
                    parts = content.split("|", 1)
                    vote_target = parts[0].strip()
                    reason = parts[1].strip()

                found = False
                for valid in players:
                    if valid in vote_target and valid != p:
                        vote_target = valid
                        found = True
                        break
                if not found:
                    others = [pl for pl in players if pl != p]
                    vote_target = random.choice(others) if others else p
                    reason = "随机投票(推理无效)"
                print(f"   🤖 [{p}] 投给了 -> {vote_target} (理由: {reason})")
            except:
                print(f"   🤖 [{p}] 弃票")
                vote_target = "弃票"
        new_votes[p] = vote_target
    print("-" * 30)
    return {"votes": new_votes}


def judge_node(state: GameState):
    """判决节点 (逻辑不变)"""
    if state.get("cheater_detected"):
        cheater = state["cheater_detected"]
        spy = state["spy_player"]
        if cheater == spy:
            return {"game_result": f"卧底[{cheater}]违规被罚下！平民胜利！"}
        else:
            return {"game_result": f"平民[{cheater}]违规被罚下！但这局算平民输。"}

    votes = state["votes"]
    counts = {}
    for target in votes.values():
        if target != "弃票": counts[target] = counts.get(target, 0) + 1
    print(f"📊 最终票数统计: {counts}")

    if not counts:
        return {"round_count": state["round_count"] + 1, "current_turn_idx": 0}

    sorted_votes = sorted(counts.items(), key=lambda x: x[1], reverse=True)
    top_votes = sorted_votes[0][1]
    top_players = [p for p, v in sorted_votes if v == top_votes]

    if len(top_players) > 1:
        print(f"⚖️ 平票！{top_players}，进入下一轮！\n")
        return {"round_count": state["round_count"] + 1, "current_turn_idx": 0}

    out_player = top_players[0]
    print(f"\n❌ [{out_player}] 被淘汰！")
    spy = state["spy_player"]
    print(
        f"🔑 真相：卧底是[{spy}]，词【{state['player_words'][spy]}】。平民词【{state['player_words'][HUMAN_NAME] if spy != HUMAN_NAME else state['player_words']['AI_1号']}】")

    if out_player == spy:
        result = f"\n🎉 平民胜利！"
    else:
        result = f"\n😈 卧底胜利！"
    return {"game_result": result}


# ---------------- 3. 构建图与路由 (核心修改) ----------------

def turn_router(state: GameState):
    """决定发言循环是否结束"""
    if state.get("cheater_detected"): return "judge"
    # 如果还有人没说过话，继续发言
    if state["current_turn_idx"] < len(state["players"]):
        return "speak"
    # 所有人都说完了，进入分析阶段
    return "analyze"


def game_router(state: GameState):
    if state.get("game_result"): return "end"
    return "continue"


workflow = StateGraph(GameState)

# 添加所有节点
workflow.add_node("Init", init_game_node)
workflow.add_node("Speaker", speaker_node)
workflow.add_node("FactCheck", fact_check_node)
# 【重点】添加分析节点
workflow.add_node("SemanticAnalysis", semantic_analysis_node)
workflow.add_node("Voting", voting_node)
workflow.add_node("Judge", judge_node)

# 编排流程
workflow.set_entry_point("Init")
workflow.add_edge("Init", "Speaker")

# Speaker -> FactCheck
workflow.add_edge("Speaker", "FactCheck")

# FactCheck -> 路由 (继续发言 OR 开始分析)
workflow.add_conditional_edges(
    "FactCheck",
    turn_router,
    {
        "speak": "Speaker",  # 继续下一位发言
        "analyze": "SemanticAnalysis",  # 【重点】所有人都说完后，去分析
        "judge": "Judge"  # 作弊直接判决
    }
)

# 【重点】分析完 -> 投票
workflow.add_edge("SemanticAnalysis", "Voting")

# 投票 -> 判决
workflow.add_edge("Voting", "Judge")

# 判决 -> 路由 (结束 OR 下一轮)
workflow.add_conditional_edges(
    "Judge",
    game_router,
    {
        "continue": "Speaker",
        "end": END
    }
)

app = workflow.compile()

# ---------------- 4. 启动游戏 ----------------
if __name__ == "__main__":
    try:
        app.invoke({"players": []})
    except KeyboardInterrupt:
        print("\n🛑 游戏已手动终止。")
    except Exception as e:
        # print(f"\n❌ 游戏发生错误: {e}")
        import traceback

        traceback.print_exc()